import gym
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from gym.utils import seeding
import random

import configparser
import tasks
from nodes import Node_U, Node_W2
import queue


# matplotlib.use("Agg")  # 控制绘图不显示，必须在import matplotlib.pyplot as plt前运行
config = configparser.ConfigParser()
config.read('env_config.ini')
new_obs_dim = 22
new_action_dim = 5
class NetworkEnv(gym.Env):
    """

    """

    metadata = {"render.modes": ["human"]}

    '''
        初始化参数设置
    '''
    def __init__(
            self,
            args,                     # 其他参数
    ):
        # 重写观测空间维度
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(new_obs_dim,))

        # 重写动作空间维度
        self.action_space = gym.spaces.Discrete(new_action_dim)

        # 初始化配置参数设置
        self.computing_node = args.NumOfWorkers
        self.consumer_node = args.NumOfUsers
        self.done_step = args.totalStep  #一个回合分配的时间步
        self.step_num = 0 #当前时间步
        self.seed = args.RamSeed
        random.seed(self.seed)
        self.fig, self.ax = plt.subplots(1, 1, figsize=(10, 7))

        ## 画图相关参数配置
        self.ax.set_ylim([0, 100])
        self.ax.set_xlabel('Step')
        self.ax.set_ylabel('Value')
        self.ax.set_title('Worker 0 State Changes')
        self.ax.plot([], [], '-b')
        self.ax.plot([], [], '-r')
        self.ax.legend(['compute', 'band'])

        plt.ion()  # 开启交互模式

        self.args = args

        self.computing_ca = [15, 20, 30]  # 算力可选列表
        self.bandwidth_ca = [9, 20, 30, 50]  # 算力带宽列表

        self.computing_re = [10]  # index=10时，用户初始化
        self.bandwidth_re = [10]  # index=10，用户初始化
        
        # 初始化固定参数设置
        self.worker_list = []       # 工作节点类列表
        self.user_list = []
        self.task_queue = queue.Queue()
        self.cl = []                # 存储状态变量 计算负载 的变化情况
        self.bl = []                # 存储状态变量  带宽负载 的变化情况


        for index in range(self.computing_node):
            self.worker_list.append(
                Node_W2(
                    self.computing_ca[index % len(self.computing_ca)],
                    self.bandwidth_ca[index % len(self.bandwidth_ca)],
                    index,
                    self.args.NodeStrategyCo,
                    self.args.NodeStrategyTr
                )
            )
        
        for index in range(self.consumer_node):
            self.user_list.append(
                Node_U(
                    index,index % 10,
                    self.computing_re[0],self.bandwidth_re[0]
                )
            )   # 通过 index % 10+1 选择不同的任务，范围 10-100 或自设值

    def step(self, action):
        # print(self.task_queue.qsize())
        task = self.task_queue.get()
        # print(task.taskID)
        self.worker_list[action].allocate_task('C', task)  #任务追加到计算节点的计算队列
        state_next = self.get_state()
        r = None # 待任务完成后计算奖励;
        done=0
        tid = task.taskID

        self.step_num+=1
        if self.step_num ==self.done_step:
            done=1
        return state_next,r,done,tid

    def reset(self):
        # 初始化固定参数设置

        self.task_queue.queue.clear() #情况任务队列
        self.cl = []  # 存储状态变量 计算负载 的变化情况
        self.bl = []  # 存储状态变量  带宽负载 的变化情况
        self.step_num = 0  # 时间步归0

        for i in range(len(self.worker_list)): #清空节点队列
            self.worker_list[i].task_process_list = []
            self.worker_list[i].task_transmit_list = []
        state = self.get_state()
        return state




    def get_state(self):
        if self.task_queue.qsize()==0:
            state =[0,0]
        else:
            task = self.task_queue.queue[0]
            state = [task.computing_re,task.bandwidth_re]  #待分配状态
        for node in self.worker_list:
            cc, bc, rc, rb = node.get_node_state()
            state.extend([cc, bc, rc, rb])

        # print(f'开始分配前，状态信息为：{state}')
        return state

    def check_idle(self): #检测是否空闲
        allrc=0
        allrb=0
        if self.task_queue.qsize() == 0:
            for node in self.worker_list:
                cc, bc, rc, rb = node.get_node_state()
                allrb +=rb
                allrc +=rc
            if allrb==0 and allrc==0:
                return True
        return False


    # 将用户产生的任务放在 task_queue 内
    def add_task(self, task):
        self.task_queue.put(task)

    # 画图函数
    def render(self, mode="human", close=False):
        self.ax.set_xlim([0, len(self.cl) if len(self.cl) > 100 else 100])  #
        self.ax.set_ylim([0, 100 if len(self.cl)==0 else max(max(self.cl),max(self.bl))])
        self.ax.lines[0].set_xdata(range(len(self.cl)))
        self.ax.lines[1].set_xdata(range(len(self.bl)))
        self.ax.lines[0].set_ydata(self.cl)
        self.ax.lines[1].set_ydata(self.bl)
        self.ax.relim()
        self.ax.autoscale_view(True, True, True)
        plt.draw()
        if mode == 'human':
            plt.pause(0.1)
        # 暂停一小段时间，允许交互窗口更新

    def _seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]




